package buff;

import java.io.File;
import java.io.IOException;
import java.io.RandomAccessFile;
import java.nio.ByteBuffer;
import java.nio.channels.FileChannel;
import java.nio.channels.FileChannel.MapMode;
import java.util.Arrays;

/**
 * License: Apache License 2.0 (http://www.apache.org/licenses/LICENSE-2.0.html)
 * <p/>
 * Author: Ashwin Jayaprakash
 * <p/>
 * Email: ashwin.jayaprakash@gmail.com
 * <p/>
 * Web: http://www.ashwinjayaprakash.com
 */
public class JavaBufferTest {
    protected static final String MSG_TEMPLATE_RESULT_HEADER =
            "test_type memory_jvm_total_bytes memory_total_buffers_bytes" +
                    " time_total_buffer_creation_millis time_total_buffer_work_millis block_size total_bytes_touched%n";

    protected static final String MSG_TEMPLATE_RESULTS = "%s %10d %10d %6d %6d %6d %10d%n";

    protected String testType;

    protected int blockSizeBytes;

    protected int numBlocks;

    public JavaBufferTest() {
    }

    public String getTestType() {
        return testType;
    }

    public void setTestType(String testType) {
        this.testType = testType;
    }

    public int getBlockSizeBytes() {
        return blockSizeBytes;
    }

    public int getNumBlocks() {
        return numBlocks;
    }

    public void setBlockSpec(int blockSizeBytes, int numBlocks) {
        this.blockSizeBytes = blockSizeBytes;
        this.numBlocks = numBlocks;
    }

    protected void beforeTestLoop() {
        System.out.printf(MSG_TEMPLATE_RESULT_HEADER);
    }

    protected void testActual(BufferMaker bufferMaker) throws Exception {
        ByteBuffer[] mbbs = new ByteBuffer[numBlocks];

        long startMaking = System.nanoTime();
        for (int i = 0; i < mbbs.length; i++) {
            mbbs[i] = bufferMaker.make(blockSizeBytes);

            if (mbbs[i].capacity() != mbbs[i].limit() || mbbs[i].capacity() != blockSizeBytes) {
                throw new RuntimeException(
                        String.format("The buffer's capacity [%d], liimt [%d] and expected size [%d] are not the same",
                                mbbs[i].capacity(), mbbs[i].limit(), blockSizeBytes));
            }
        }
        long endMaking = System.nanoTime();

        long startPounding = System.nanoTime();
        long totalBytesTouched = 0;
        for (ByteBuffer mbb : mbbs) {
            int bytesTouched = poundBuffer(mbb);

            if (bytesTouched != blockSizeBytes || mbb.remaining() != 0) {
                throw new RuntimeException(
                        "Not all bytes in the buffer were touched. Only [" + bytesTouched + "]" +
                                " were affected. Expected [" + blockSizeBytes + "]");
            }

            totalBytesTouched += bytesTouched;
        }
        if (totalBytesTouched != (numBlocks * blockSizeBytes)) {
            throw new RuntimeException(
                    "Not all bytes in all the buffer spaces were touched." +
                            " Only [" + totalBytesTouched + "]" +
                            " were affected. Expected [" + (numBlocks * blockSizeBytes) + "]");
        }
        long endPounding = System.nanoTime();

        clearAndPrint(mbbs, (endMaking - startMaking), (endPounding - startPounding), totalBytesTouched);
    }

    protected void afterTestLoop() {
        System.out.println();
    }

    /**
     * @param mbb
     * @return The total number of bytes that were touched.
     */
    protected int poundBuffer(ByteBuffer mbb) {
        int limit = mbb.limit();

        int bytesWritten = 0;
        for (; bytesWritten < limit; ) {
            mbb.putInt(bytesWritten);
            bytesWritten += 4;
        }

        return bytesWritten;
    }

    protected void clearAndPrint(ByteBuffer[] mbbs, long bufferCreationTimeNanos, long bufferWorkTimeNanos,
                                 long totalBytesTouched) {
        Runtime runtime = Runtime.getRuntime();
        long freeBefore = 0;
        long totalAfter = 0;
        long freeAfter = 0;

        System.gc();
        System.gc();
        freeBefore = runtime.freeMemory();
        for (int i = 0, x = 0; i < mbbs.length; i++) {
            mbbs[i] = null;
        }
        System.gc();
        System.gc();
        totalAfter = runtime.totalMemory();
        freeAfter = runtime.freeMemory();

        System.out.printf(MSG_TEMPLATE_RESULTS, testType,
                totalAfter, (freeAfter - freeBefore),
                bufferCreationTimeNanos / 1000000,
                bufferWorkTimeNanos / 1000000, blockSizeBytes, totalBytesTouched);
    }

    protected void testDisk(File file, String mode) throws Exception {
        if (file.exists() && !file.delete()) {
            throw new RuntimeException(
                    "Unable to delete file [" + file.getAbsolutePath() + "] for new round of test");
        }
        file.createNewFile();

        RandomAccessFile raf = new RandomAccessFile(file, mode);
        raf.setLength(numBlocks * blockSizeBytes);

        try {
            final FileChannel fc = raf.getChannel();

            BufferMaker bufferMaker = new BufferMaker() {
                int startPosition = 0;

                @Override
                public ByteBuffer make(int size) throws IOException {
                    ByteBuffer bb = fc.map(MapMode.READ_WRITE, startPosition, size);
                    startPosition += size;

                    return bb;
                }
            };

            testActual(bufferMaker);
        }
        finally {
            raf.close();
        }
    }

    protected void testMemory(final boolean direct) throws Exception {
        BufferMaker bufferMaker = new BufferMaker() {
            @Override
            public ByteBuffer make(int size) throws IOException {
                return direct ? ByteBuffer.allocateDirect(size) : ByteBuffer.allocate(size);
            }
        };

        testActual(bufferMaker);
    }

    protected boolean runTest(String[] args, int loopCount) throws Exception {
        int blockSizeBytes = 0;
        int numBlocks = 0;

        if (args[0].equalsIgnoreCase("heap")) {
            setTestType(args[0]);
            blockSizeBytes = Integer.parseInt(args[1]);
            numBlocks = Integer.parseInt(args[2]);
            setBlockSpec(blockSizeBytes, numBlocks);

            beforeTestLoop();
            for (int i = 0; i < loopCount; i++) {
                testMemory(false);
            }
            afterTestLoop();
        }
        else if (args[0].equalsIgnoreCase("direct")) {
            setTestType(args[0]);
            blockSizeBytes = Integer.parseInt(args[1]);
            numBlocks = Integer.parseInt(args[2]);
            setBlockSpec(blockSizeBytes, numBlocks);

            beforeTestLoop();
            for (int i = 0; i < loopCount; i++) {
                testMemory(true);
            }
            afterTestLoop();
        }
        else if (args[0].equalsIgnoreCase("file")) {
            File f = new File(args[1]);
            String mode = args[2];
            blockSizeBytes = Integer.parseInt(args[3]);
            numBlocks = Integer.parseInt(args[4]);
            setTestType(args[0] + "_" + args[1] + "_" + args[2]);
            setBlockSpec(blockSizeBytes, numBlocks);

            beforeTestLoop();
            for (int i = 0; i < loopCount; i++) {
                testDisk(f, mode);
            }
            afterTestLoop();
        }
        else {
            System.err.println("Unrecognized parameters " + Arrays.asList(args));
            System.err.println("Options:");
            System.err.println("    heap <block_size_bytes> <num_blocks>");
            System.err.println("    direct <block_size_bytes> <num_blocks>");
            System.err.println("    file <path> <rw|rws|rwd> <block_size_bytes> <num_blocks>");

            return false;
        }

        return true;
    }

    public static void main(String[] args) throws Exception {
        new JavaBufferTest().runTest(args, 24);
    }

    public static interface BufferMaker {
        ByteBuffer make(int size) throws Exception;
    }
}
